import cv2
import time
import os
from datetime import datetime
from ultralytics import YOLOv10
from qcloud_cos import CosConfig
from qcloud_cos import CosS3Client
import logging
import sys
import threading

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s: %(message)s')

class VideoDetector:
    def __init__(self):
        # 初始化参数
        self.image_size = 640
        self.conf_threshold = 0.1
        self.model = YOLOv10("models/YOLOv10-FireSmoke-M.pt")
        
        # 文件路径设置
        self.output_dir = "detected_frames"
        os.makedirs(self.output_dir, exist_ok=True)
        
        # COS配置
        self.cos_bucket = "hujiaqi-1358073331"
        self.cos_folder = "fire_detections/"
        
        # 上传间隔(秒)
        self.save_interval = 5
        self.last_save_time = 0
        
        # 启动上传线程
        self.upload_thread = threading.Thread(target=self.monitor_and_upload)
        self.upload_thread.daemon = True
        self.upload_thread.start()

    def detect_frame(self, frame):
        """执行目标检测"""
        results = self.model.predict(source=frame, imgsz=self.image_size, conf=self.conf_threshold)
        return results[0].plot(), len(results[0].boxes) > 0

    def save_detection(self, frame):
        """保存检测结果"""
        current_time = time.time()
        if current_time - self.last_save_time >= self.save_interval:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"detection_{timestamp}.jpg"
            local_path = os.path.join(self.output_dir, filename)
            cv2.imwrite(local_path, frame)
            self.last_save_time = current_time
            logging.info(f"检测结果已保存: {local_path}")
            return True
        return False

    def upload_to_cos(self, local_path):
        """上传文件到COS并删除本地文件"""
        try:
            # 建议从环境变量获取，这里简化处理
            secret_id = 'AKIDuuiJhsqa9UbUWbklojfYDfvxlNn2A8tQ'
            secret_key = 'Ogw7BEvSnvpUwHWkrwwnStENI0UnvqY2'
            region = 'ap-guangzhou'
            
            config = CosConfig(Region=region, SecretId=secret_id, SecretKey=secret_key)
            client = CosS3Client(config)
            
            if not os.path.exists(local_path):
                logging.warning(f"文件不存在: {local_path}")
                return False
                
            # 上传到COS
            cos_path = self.cos_folder + os.path.basename(local_path)
            client.put_object_from_local_file(
                Bucket=self.cos_bucket,
                LocalFilePath=local_path,
                Key=cos_path
            )
            logging.info(f"已上传到COS: {cos_path}")
            
            # 上传成功后删除本地文件
            os.remove(local_path)
            logging.info(f"已删除本地文件: {local_path}")
            return True
            
        except Exception as e:
            logging.error(f"上传失败: {str(e)}")
            return False

    def monitor_and_upload(self):
        """监控文件夹并上传新文件"""
        while True:
            try:
                # 获取文件夹内所有jpg文件
                files = [f for f in os.listdir(self.output_dir) if f.endswith('.jpg')]
                
                for filename in files:
                    filepath = os.path.join(self.output_dir, filename)
                    if self.upload_to_cos(filepath):
                        time.sleep(0.5)  # 避免频繁上传
                        
                time.sleep(1)  # 每1秒检查一次
                
            except Exception as e:
                logging.error(f"上传线程错误: {str(e)}")
                time.sleep(5)

    def process_video(self, video_path):
        """处理视频文件"""
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            logging.error(f"无法打开视频文件: {video_path}")
            return
            
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        processed_frames = 0
        
        logging.info(f"开始处理视频: {video_path}")
        logging.info(f"总帧数: {total_frames}")
        logging.info("按Q键可提前终止...")

        while True:
            ret, frame = cap.read()
            if not ret:
                break
                
            processed_frames += 1
            
            # 执行检测
            processed_frame, has_detection = self.detect_frame(frame)
            
            # 显示处理进度
            if processed_frames % 10 == 0:
                logging.info(f"处理进度: {processed_frames}/{total_frames} 帧")
            
            # 保存检测结果
            if has_detection:
                self.save_detection(processed_frame)
            
            # 显示实时画面
            cv2.imshow("YOLOv10 实时检测", processed_frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
                
        cap.release()
        cv2.destroyAllWindows()
        logging.info("视频处理完成！")

if __name__ == "__main__":
    detector = VideoDetector()
    detector.process_video("1.mp4")